import os
# 解决OMP冲突: 必须在导入numpy或torch之前设置
os.environ['KMP_DUPLICATE_LIB_OK']='True'

import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm

from model import SimpleCNN
from data_loader import get_data_loaders

# --- 配置参数 ---
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DATA_DIR = './data'
MODEL_SAVE_PATH = './gtsrb_cnn.pth'
NUM_EPOCHS = 15
LEARNING_RATE = 0.001
BATCH_SIZE = 64
NUM_CLASSES = 43

def train_model():
    """
    训练并保存交通标志识别模型。
    """
    print(f"使用设备: {DEVICE}")

    # 1. 加载数据
    try:
        train_loader, test_loader = get_data_loaders(DATA_DIR, BATCH_SIZE)
        print("数据加载成功。")
    except Exception as e:
        print(f"数据加载失败: {e}")
        print(f"请确保Kaggle数据集已正确解压到 '{os.path.abspath(DATA_DIR)}' 目录下。")
        return

    # 2. 初始化模型、损失函数和优化器
    model = SimpleCNN(num_classes=NUM_CLASSES).to(DEVICE)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
    
    # 用于记录最佳模型
    best_accuracy = 0.0

    print("开始训练...")
    # 3. 训练循环
    for epoch in range(NUM_EPOCHS):
        # --- 训练阶段 ---
        model.train()
        running_loss = 0.0
        train_loop = tqdm(train_loader, desc=f"Epoch {epoch+1}/{NUM_EPOCHS} [训练]")
        
        for images, labels in train_loop:
            images, labels = images.to(DEVICE), labels.to(DEVICE)
            
            # 梯度清零
            optimizer.zero_grad()
            
            # 前向传播
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            # 反向传播和优化
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item() * images.size(0)
            train_loop.set_postfix(loss=loss.item())

        epoch_loss = running_loss / len(train_loader.dataset)
        
        # --- 验证阶段 ---
        model.eval()
        correct = 0
        total = 0
        val_loss = 0.0
        
        val_loop = tqdm(test_loader, desc=f"Epoch {epoch+1}/{NUM_EPOCHS} [验证]")
        with torch.no_grad():
            for images, labels in val_loop:
                images, labels = images.to(DEVICE), labels.to(DEVICE)
                
                outputs = model(images)
                loss = criterion(outputs, labels)
                val_loss += loss.item() * images.size(0)
                
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        
        accuracy = 100 * correct / total
        val_loss /= len(test_loader.dataset)
        
        print(f"Epoch {epoch+1}/{NUM_EPOCHS} | "
              f"训练损失: {epoch_loss:.4f} | "
              f"验证损失: {val_loss:.4f} | "
              f"验证准确率: {accuracy:.2f}%")

        # 保存最佳模型
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            torch.save(model.state_dict(), MODEL_SAVE_PATH)
            print(f"检测到更优模型，已保存到 '{MODEL_SAVE_PATH}' (准确率: {best_accuracy:.2f}%)")

    print("\n训练完成！")
    print(f"最佳模型已保存在 '{MODEL_SAVE_PATH}'，最高验证准确率为: {best_accuracy:.2f}%")


if __name__ == '__main__':
    train_model() 